import cv2
import numpy as np
import pyautogui
from ultralytics import YOLO
import time

def capture_screen():
    """捕获游戏区域屏幕截图"""
    try:
        # 指定窗口区域 (0, 87, 842, 570)
        screenshot = pyautogui.screenshot(region=(0, 87, 842, 570))
        # 转换为OpenCV格式
        frame = cv2.cvtColor(np.array(screenshot), cv2.COLOR_RGB2BGR)
        return frame
    except Exception as e:
        print(f"截图错误: {e}")
        return None

def main():
    # 加载YOLOv8模型
    model = YOLO('best.pt')
    
    # 初始化FPS计算
    last_time = time.time()
    frame_count = 0
    fps = 0
    
    try:
        while True:
            # 捕获屏幕
            frame = capture_screen()
            if frame is None:
                continue

            # 使用YOLOv8进行检测
            results = model(frame, verbose=False)
            
            # 计算FPS
            frame_count += 1
            current_time = time.time()
            if current_time - last_time >= 1.0:
                fps = frame_count / (current_time - last_time)
                frame_count = 0
                last_time = current_time

            # 在图像上绘制检测结果
            for r in results:
                # 获取检测框
                boxes = r.boxes
                for box in boxes:
                    # 获取类别和置信度
                    class_id = r.names[int(box.cls)]
                    confidence = float(box.conf)
                    
                    # 只显示置信度大于0.5的检测结果
                    if confidence > 0.5:
                        # 获取边界框坐标
                        x1, y1, x2, y2 = box.xyxy[0].cpu().numpy().astype(int)
                        
                        # 根据类别选择颜色
                        if class_id == '岩浆':
                            color = (0, 0, 255)  # 红色
                        elif class_id == '石块':
                            color = (128, 128, 128)  # 灰色
                        elif class_id == '自己':
                            color = (255, 0, 0)  # 蓝色
                        elif class_id == '怪兽':
                            color = (0, 255, 255)  # 黄色
                        elif class_id == '矿':
                            color = (0, 255, 0)  # 绿色
                        else:
                            color = (255, 255, 255)  # 白色
                        
                        # 绘制边界框
                        cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
                        
                        # 添加标签
                        label = f"{class_id} {confidence:.2f}"
                        cv2.putText(frame, label, (x1, y1 - 10),
                                  cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)

            # 显示FPS
            cv2.putText(frame, f"FPS: {fps:.1f}", (10, 30),
                       cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)

            # 显示结果
            cv2.imshow('YOLOv8 Detection', frame)
            
            # 按'q'退出
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break

    except KeyboardInterrupt:
        print("程序被用户中断")
    finally:
        cv2.destroyAllWindows()

if __name__ == "__main__":
    main()
